package SSF.OS.NetFlow.Util;

import java.util.*;

/**
 * Given one flow data file, divide it into several files according to
 * the nhi address prefixes of the hosts that dump the flow data.
 * In this version, it just takes into account the fist 3 levels of nhi ID
 */
public class NHIDemuxer extends FilterPlayer
{
    private Hashtable recorderTable = null;
    public int level=1;
    public String srcFileName=null;
    protected int openFileMAX = 20;
    protected int numFileOpened = 0;
    private Vector typeVec;

    public NHIDemuxer(String streamID)
    {
	super(streamID);
	recorderTable = new Hashtable();
	typeVec = new Vector();
    }

    public NHIDemuxer(String streamID, int nLevel)
    {
	super(streamID);
	recorderTable = new Hashtable();
	typeVec = new Vector();
	level = nLevel;
    }

    @Override
	protected void addType(String spec) throws streamException
    {
	//for debug
	//System.err.println("addType");

	super.addType(spec);
	typeVec.add(spec);

	//write it to all the currently opened recorders.
	Object obj;
	for (Enumeration enum = recorderTable.elements();
	     enum.hasMoreElements();) {
	    obj = enum.nextElement();
	    if (obj instanceof BasicRecorder)
		((BasicRecorder)obj).send(0,0,0.0,spec.getBytes(),
					  0,spec.length());
	}
    }

    @Override
	protected void addSource(String spec) throws streamException
    {
	//for debug
	//System.err.println("addSource");

	int sp = spec.indexOf(" ");
	String N;
	if (sp > 0) {
	    Integer I = new Integer(spec.substring(0,sp));
	    N = spec.substring(sp+1).trim();
	    sourceDictionary.put(I,N);
	    sourceDictionary.put(N,I);
	}
	else throw new streamException("Bad code: "+spec);

	// now N is the nhi address, get its prefix and thus the recorder.
	BasicRecorder rec = getRecorder(getNhiPrefix(N, level));	

	//write in the current recorder.
	rec.send(1, 0, 0.0, spec.getBytes(), 0, spec.length());
    }

    private BasicRecorder getRecorder(String key) throws streamException
    {
	//for debug
	//System.err.println("getRecorder");

	Object obj = recorderTable.get(key);
	BasicRecorder rec = null;
	int typeVecSize = typeVec.size();
	int initTypeNum=0;
	String fname;

	if (obj == null) {
	    //this is the first time this nhi prefix is read.
	    //check how many files have been  opened in the total, 
	    //close some if necesssary.
	    checkRecTable();

	    //create a new recorder for it.
	    rec = new BasicRecorder(streamID);

	    //create a new file name
	    fname = srcFileName+"."+key;
	    rec.connectWrite(fname);
	    numFileOpened++;

	    recorderTable.put(key, rec);

	} else if (obj instanceof Integer) {
	    initTypeNum = ((Integer)obj).intValue();
	    
	    //it exists, but has been disconnected ...
	    checkRecTable();

	    //re-connect it using append mode.
	    rec = new BasicRecorder(streamID);
	    fname = srcFileName+"."+key;
	    rec.connectWrite(fname, true);
	    numFileOpened++;

	    //replace the Integer
	    recorderTable.remove(key);
	    recorderTable.put(key, rec);
	} else 
	    return (BasicRecorder)obj;

	//write in *new* type records.
	for (int i=initTypeNum; i<typeVecSize; i++)
	    rec.send(0,0,0.0,((String)(typeVec.get(i))).getBytes(),
		     0,((String)typeVec.get(i)).length());

	return rec;
    }
    @Override
	public int receive(int type_code, int source_code, double timestamp,
		       byte[] buf, int offset, int length)
    {
	//for debug
	//System.err.println("receive");

	// get the key
	String key = getNhiPrefix(getRecordSourceString(source_code),level);

	// try to get the recorder and distribute the data
	try {
	    BasicRecorder rec = getRecorder(key);
	    rec.send(type_code, source_code, timestamp, buf, offset, length);
	} catch (streamException se) {
	    System.err.println(se.getMessage());
	    se.printStackTrace();
	    System.exit(-1);
	}
	return 0;
    }

    /**
     * Given the nhi address, return its first *level* prefix.
     * E.g. the first 2 level prefix of 3:2:5:2(0) is 3:2.
     * If the level is too big, return the nhi address without the interface.
     */
    public static String getNhiPrefix(String nhi, int level)
    {
	//for debug
	//System.err.println("getNhiPrefix "+nhi+" "+level);

	int current = -1;
	int len = nhi.length();

	for (int i=0; i<level; i++) {
	    current = nhi.indexOf(':',current+1);
	    if (current == len - 1) {
		//the level is too big.
		//get rid of the interface part
		current = nhi.indexOf('(');
		break;
	    }
	}
	//finished the cycle, got the index of the prefix + 1
	if (current < 0) {
	    //should return the whole string, there is no '('.
	    return nhi;
	}
	return new String(nhi.getBytes(), 0, current);
    }

    /**
     * Close (disconnect) files when too many files are open.
     */
    private void checkRecTable()
    {
	//for debug
	//System.err.println("checkRecTable");

	Object obj = null;
	Object key;
	Integer currentTypeSize = new Integer(typeVec.size());

	if (numFileOpened < openFileMAX)
	    return;
	    
	//too many files, close the fist 10.

	for (Enumeration enum=recorderTable.keys();
	     enum.hasMoreElements();){
	    key = enum.nextElement();
	    obj = recorderTable.get(key);
	    if (obj instanceof BasicRecorder){
		((BasicRecorder)obj).disconnect();
		numFileOpened--;
		//try to replace the recorder with the size of typeVec now.
		//!!!!!COULD HAVE PROBLEM !!!!!
		recorderTable.put(key, currentTypeSize);

		if (numFileOpened == 0 || numFileOpened == openFileMAX-10)
		    break;
	    }
	}
    }

    @Override
	public void disconnect() {
	//disconnect or the recorders, clean the table.
	Object obj;
	for (Enumeration enum=recorderTable.elements();
	     enum.hasMoreElements();) {
	    obj = enum.nextElement();
	    if (obj instanceof BasicRecorder)
		((BasicRecorder)obj).disconnect();
	}
	recorderTable.clear();
	super.disconnect();
    }

    public static void main(String[] args)
    {
	String srcFile = null;
	String streamID = null;
	int nLevel = 1;

	try {
	    switch (args.length) {
	    case 4:
		if (!(args[0].equals("-l")))
		    errorUsage();
		nLevel = Integer.valueOf(args[1]).intValue();
		srcFile = args[2];
		streamID = args[3];
		break;
	    case 2:
		srcFile = args[0];
		streamID = args[1];
		break;
	    default:
		errorUsage();
	    }
	} catch (Exception any) {
	    System.err.println(any.getMessage());
	    errorUsage();
	}

	NHIDemuxer nhiDemuxer = new NHIDemuxer(streamID, nLevel);
	try {
	    nhiDemuxer.srcFileName=srcFile;
	    nhiDemuxer.connectRead(srcFile);
	    nhiDemuxer.disconnect();
	} catch (streamException se) {
	    System.err.println(se.getMessage());
	    System.exit(-1);
	}
    }

    public static void errorUsage()
    {
	System.err.println("usage: java NHIDemuxer [-l level] <netflow data file> <stream name>");
	System.exit(1);
    }
}
